#!/usr/bin/env python3

import hybrid_astar as ha
import MPC_XY_Frame as mxf
import math
import rospy
from geometry_msgs.msg import Twist
import time
import sys
import matplotlib.pyplot as plt
import copy
from mvccep_image_localization.msg import base_localization

robot_num = "robot_2"

## 装换坐标：原点中心从中心到左下，并进行缩放
def aStarPretreatment(XYY_START,XYY_END):
    XYY_START[0] = XYY_START[0] / ha.C.N + ha.C.MAP_SIZE[0] / 2
    XYY_START[1] = XYY_START[1] / ha.C.N + ha.C.MAP_SIZE[1] / 2
    XYY_END[0] = XYY_END[0] / ha.C.N + ha.C.MAP_SIZE[0] / 2
    XYY_END[1] = XYY_END[1] / ha.C.N + ha.C.MAP_SIZE[1] / 2


## a星轨迹后处理：将坐标转换为原来
def aStarPostProcessing(path):
    ax,ay = [],[]
    for i in range(len(path.x)):
        ax.append((path.x[i] - ha.C.MAP_SIZE[0]/2) * ha.C.N)
        ay.append((path.y[i] - ha.C.MAP_SIZE[1]/2) * ha.C.N)
    return ax,ay

# 四元数计算车辆姿态角
def eulerAnglesToEulerAngles(q):
    siny_cosp = 2 * (q.w * q.z + q.x * q.y)
    cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
    yaw = math.atan2(siny_cosp, cosy_cosp)

    return yaw

# 计数速度
def getV(x1,y1,t1,x2,y2,t2):
    return math.sqrt((x1-x2)**2 + (y1-y2)**2)/(t2-t1)

# 更新车辆位置信息
def updateNodeFromSim(Node,data):
    x_cur = data.robot_2.pos.x
    y_cur = data.robot_2.pos.x
    time_cur = round(time.time(),2)
    node.v = getV(node.x,node.y,node.time,x_cur,y_cur,time_cur)
    node.x = x_cur
    node.y = y_cur
    node.yaw = data.robot_2.angle
    node.time = time_cur
    

def callback(data):
    ## MPC规划
    startTime = time.time()
    global a_opt, delta_opt, x_opt, y_opt, yaw_opt, v_opt

    if isActual:
        updateNodeFromSim(node,data)    # 更新小车状态

    z_ref, target_ind = mxf.calc_ref_trajectory_in_T_step(node, ref_path, sp)   # 计算预测结果
    z0 = [node.x, node.y, node.v, node.yaw]     
    a_opt, delta_opt, x_opt, y_opt, yaw_opt, v_opt = mxf.linear_mpc_control(z_ref, z0, a_opt, delta_opt)    # MPC

    ###################
    if delta_opt is not None:
        delta_exc,a_exc = delta_opt[0],a_opt[0]

    linear_x = 0
    if isActual:
        node_new = copy.deepcopy(node)
        node_new.update(a_exc,delta_exc,1.0)
        linear_x = node_new.v
    else:
        node.update(a_exc,delta_exc,1.0)
        linear_x = node.v
    ###################


    ## 发布控制量
    # 计算控制量
    # if delta_opt is not None:
    #     v_exc = v_opt[0] + a_opt[0]*dt      # cm/s
    #     delta_exc = delta_opt[0]            # 倾角 rad

    # 发布控制量
    if isActual:
        twist.linear.x = linear_x               # 线速度 m/s
        twist.angular.z = delta_exc             # 倾角 rad
        cmd_pub.publish(twist)


    ## 终点判断
    dist = math.hypot(node.x - cx[-1], node.y - cy[-1])             # 计算终点距离
    print("终点距离:",dist)
    print("速度:",node.v,"倾角:",delta_exc)
    print("小车信息,x:",node.x,"y:",node.y,"v:",node.v,"yaw:",math.degrees(node.yaw))
    if dist < P.dist_stop and node.v < mxf.P.speed_stop:           # 终点判断
        plt.show()
        sys.exit()
    endTime = time.time()
    print("speed time:",round(endTime - startTime,3))


    # 控制发布信息评率
    timeSleep = dt - (endTime - startTime)
    # print("timeSleep:",timeSleep)
    if timeSleep > 0:
        time.sleep(timeSleep)

    #########################
    if isPlot:
        yaw.append(node.yaw)
        x.append(node.x)
        y.append(node.y)
        dy = (node.yaw - yaw[-2]) / (node.v * P.dt)
        steer = mxf.rs.pi_2_pi(-math.atan(P.WB * dy)) # 计算前轮倾角
        plt.cla()
        mxf.draw.draw_car(node.x, node.y, node.yaw, steer, P)
        plt.gcf().canvas.mpl_connect('key_release_event',
        lambda event:
        [exit(0) if event.key == 'escape' else None])

        if x_opt is not None:
            plt.plot(x_opt, y_opt, color='darkviolet', marker='*')
            plt.plot(cx, cy, color='gray')
            plt.plot(x, y, '-b')
            plt.plot(cx[target_ind], cy[target_ind])
            plt.axis("equal")
            plt.title("Linear MPC, " + "v = " + str(round(node.v * 3.6, 2)))
            plt.pause(0.001)
    #########################


if __name__ == '__main__':
    global isActual,isPlot
    isActual = True
    isPlot = False

    ## a星轨迹规划
    global dt,cmd_pub,twist,node,sp
    dt = mxf.P.dt
    XYY_START = [0, 0.0, math.radians(45)]    # [m] 开始位置
    XYY_END = [0.8, 0.8, math.radians(90)]   # [m] 结束位
    aStarPretreatment(XYY_START,XYY_END)    # 数据预处理
    path = ha.main(XYY_START,XYY_END,True)   # a星规划
    ax,ay = aStarPostProcessing(path)   # 数据后处理


    ## MPC轨迹控制初始化
    P = mxf.P
    cx, cy, cyaw, ck, s = mxf.cs.calc_spline_course(ax, ay, ds=P.d_dist)    # 生成参考轨迹
    sp = mxf.calc_speed_profile(cx, cy, cyaw, P.target_speed)   # 初略速度规划
    ref_path = mxf.PATH(cx, cy, cyaw, ck)                   # 参考轨迹
    node = mxf.Node(x=ax[0], y=ay[0], yaw=cyaw[0], v=0.0,time=round(time.time(),2))   # 当前点信息
    global a_opt, delta_opt, x_opt, y_opt, yaw_opt, v_opt,yaw,x,y
    delta_opt,a_opt = None,None
    x,y,yaw = [node.x],[node.y],[node.yaw]
    

    # 节点初始化
    if isActual:
        rospy.init_node("one_robot_astar_mpc",anonymous=False)  # 节点
        cmd_pub = rospy.Publisher("/" + robot_num + "/cmd_vel" ,Twist, queue_size=1)  # 发布者
        rospy.Subscriber("base_localization" ,base_localization, callback,queue_size=1) # 接受者
        twist = Twist()     # 发布对象:速度类型
        rospy.spin()

    ###################
    while 1:
        callback(1)
    ###################